#include <TlHelp32.h>
#include <iostream>
#include <psapi.h>

#pragma comment(lib, "psapi")
//#pragma once
namespace Injection{
			DWORD getProcessId(LPSTR szExeName);
			DWORD FindModule(LPSTR szModuleName);
			DWORD ListModules(LPSTR szProcessName);
			DWORD ListProcesses();
			DWORD KillProcess(LPSTR szProcessName);
			DWORD AttachDll(LPSTR szModuleName, LPSTR szProcessName);

			DWORD ListProcesses()
			{
				int b =0;
				PROCESSENTRY32 Pc = { sizeof(PROCESSENTRY32) };
				HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
				if(Process32First(hSnapshot, &Pc)){
						do{
								b++;
								cout << Pc.szExeFile << endl;
						}while(Process32Next(hSnapshot, &Pc));
						cout << b << " running processes";
				}
				CloseHandle(hSnapshot);
				return NULL;
			}

			DWORD KillProcess(const char* szProcessName)
			{
					DWORD ExitCode = NULL;
					HANDLE hProc;
					PROCESSENTRY32 Pc = { sizeof(PROCESSENTRY32) };
					HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
					if(Process32First(hSnapshot, &Pc)){
							do{
								if(strcmp((const char*)szProcessName, (const char*)Pc.szExeFile)==0)
									{
											hProc = OpenProcess(PROCESS_ALL_ACCESS, false, Pc.th32ProcessID);
											GetExitCodeProcess(hProc, &ExitCode);
											TerminateProcess(hProc, ExitCode);
											CloseHandle(hProc);
									}
							}while(Process32Next(hSnapshot, &Pc));
					}
					CloseHandle(hSnapshot);
					return NULL;
			}

			DWORD ListModules(LPSTR szProcessName)
			{
					char DllName[MAX_PATH] = {NULL};
					HANDLE hProc;
					DWORD  dwProcessId;
					DWORD Needed = NULL;
					HMODULE hMod[1024];


					cout << "listing modules of " << szProcessName << "..." << endl;

					if((dwProcessId = getProcessId((LPSTR)szProcessName)) == NULL) {
							cout << "Couldn't find " << szProcessName;
							return (0);
					}

					hProc = OpenProcess(PROCESS_ALL_ACCESS, false, dwProcessId);
					EnumProcessModules(hProc, hMod, sizeof(hMod), &Needed);
					for (int i=0;i<(Needed / sizeof(HMODULE)); i++ )
							{                       
									if(GetModuleFileNameExA(hProc, hMod[i], DllName, sizeof(DllName)/sizeof(char)))
									{
											cout << DllName << endl;
									}
							}
					CloseHandle(hProc);
					return NULL;
			}

			DWORD getProcessId(LPSTR szExeName) {
					PROCESSENTRY32 Pc = { sizeof(PROCESSENTRY32) };
					HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
					if(Process32First(hSnapshot, &Pc)){
							do{
									if(!strcmp((const char*)Pc.szExeFile, szExeName)) {
											return Pc.th32ProcessID;
									}
							}while(Process32Next(hSnapshot, &Pc));
					}
					CloseHandle(hSnapshot);
					return NULL;
			}

			DWORD FindModule(LPSTR szModuleName) {
					PROCESSENTRY32 Pc = { sizeof(PROCESSENTRY32) };
					HMODULE hMod[1024];
					DWORD Needed = NULL;
					HANDLE hProc;
					char DllName[MAX_PATH] = {NULL};
					int b=0;

					HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
					if(Process32First(hSnapshot, &Pc)){
							do{
									hProc = OpenProcess(PROCESS_ALL_ACCESS, false, Pc.th32ProcessID);
									EnumProcessModules(hProc, hMod, sizeof(hMod), &Needed);
									for (int i=0;i<(Needed / sizeof(HMODULE)); i++ )
											{                       
													if(GetModuleFileNameEx(hProc, hMod[i], (LPWSTR)DllName, sizeof(DllName)/sizeof(char)))
													{
															if(strstr(DllName, szModuleName)!=NULL)
															{
																	cout << "Found Dll in " << Pc.szExeFile << endl << DllName << endl;
																	b++;
															}
													}
											}
							}while(Process32Next(hSnapshot, &Pc));
							cout << "Found dll " << b << " times";
					}       
					CloseHandle(hProc);
					CloseHandle(hSnapshot);
					return NULL;
			}

			DWORD AttachDll(LPSTR szModuleName, LPSTR szProcessName)
			{
					DWORD dwProcessId = getProcessId(szProcessName);
					DWORD ExitCode = NULL;
					HANDLE hProc;
					HANDLE thread;
					hProc = OpenProcess(PROCESS_ALL_ACCESS, false, dwProcessId);
					void *AllocSpace = VirtualAllocEx(hProc, NULL,strlen(szModuleName)+1, MEM_COMMIT, PAGE_READWRITE);
					DWORD error = GetLastError();
					BOOL status = WriteProcessMemory(hProc, AllocSpace, (void*)szModuleName, strlen(szModuleName)+1,NULL );
					error = GetLastError();
					thread = CreateRemoteThread(hProc, NULL, 0,(LPTHREAD_START_ROUTINE )GetProcAddress(GetModuleHandle((LPCWSTR)"Kernel32"), "LoadLibraryA"),AllocSpace, 0, NULL );
					error = GetLastError();
					WaitForSingleObject(thread, INFINITE);
					GetExitCodeThread(thread, &ExitCode);
					VirtualFreeEx(hProc, AllocSpace, sizeof(szModuleName), MEM_RELEASE);
					CloseHandle(thread);
					CloseHandle(hProc);
					return NULL;
			}
}